import mlflow
import pandas as pd

def generate_recommendations_with_comparision(
    experiment_ids,
    aggregation_function="common_features",
    main_note="sizes_acts",
    note_to_compare="sizes_L2_without_acts",
    group_type="sim"
):
    all_rows = []

    for exp_id in experiment_ids:
        runs = mlflow.search_runs(
            experiment_ids=[exp_id],
            output_format="list"
        )

        for run in runs:
            note = run.data.params.get("note")
            fusion = run.data.params.get("SAE_fusion_strategy")
            group = run.data.params.get("group_type")

            if fusion != aggregation_function or group != group_type:
                continue

            dataset = run.data.params.get("dataset", f"Exp-{exp_id}")
            dim = int(run.data.params.get("embedding_dim", 0))
            topk = int(run.data.params.get("top_k", 0))

            row_key = (dataset, dim, topk)
            model_type = "main" if note == main_note else "compare" if note == note_to_compare else None
            if not model_type:
                continue

            metrics = {
                ("G/mean", model_type): run.data.metrics.get("CommonItemsNDCG20/median"),
                ("U/mean", model_type): run.data.metrics.get("NDCG20/mean"),
                ("U/min", model_type): run.data.metrics.get("NDCG20/min"),
                ("Pop", model_type): run.data.metrics.get("Popularity/mean"),
            }

            all_rows.append((row_key, metrics))

    # Combine metrics into a dictionary
    records = {}
    for key, metrics in all_rows:
        if key not in records:
            records[key] = {}
        records[key].update(metrics)

    df = pd.DataFrame.from_dict(records, orient="index")
    df.index.names = ["Dataset", "Dimensions", "TopK"]

    # Calculate % differences
    result_cols = []
    for metric in sorted(set(k[0] for k in df.columns)):
        main_col = (metric, "main")
        compare_col = (metric, "compare")
        percent_col = (metric, "% change")

        if main_col in df.columns and compare_col in df.columns:
            df[percent_col] = ((df[main_col] - df[compare_col]) / df[compare_col].abs()) * 100
            result_cols.extend([main_col, percent_col, compare_col])
        elif main_col in df.columns:
            result_cols.append(main_col)

    # Keep only main and percent change columns, sort them by metric
    df = df[result_cols]
    df = df.round(2)
    df = df.sort_index(axis=1, level=0).sort_values(by=["Dataset", "Dimensions", "TopK"])
    

    return df.reset_index()

import mlflow
import pandas as pd

def generate_sae_table_with_comparison(experiment_ids, main_note, note_to_compare):
    all_rows = []

    for exp_id in experiment_ids:
        runs = mlflow.search_runs(
            experiment_ids=[exp_id],
            output_format="list"
        )

        for run in runs:
            note = run.data.params.get("note")
            if note not in [main_note, note_to_compare]:
                continue

            dataset = run.data.params.get("dataset", f"Exp-{exp_id}")
            dim = int(run.data.params.get("embedding_dim", 0))
            base_factors = float(run.data.params.get("base_factors", 0.0))
            expansion_ratio = dim / base_factors if base_factors else None
            topk = int(run.data.params.get("top_k", 0))

            row_key = (dim, round(expansion_ratio, 3), topk)
            model_type = "main" if note == main_note else "compare"

            metrics = {
                (dataset, model_type, "CS"): run.data.metrics.get("CosineSim/test"),
                (dataset, model_type, "Deg"): run.data.metrics.get("NDCG20_Degradation/test"),
                (dataset, model_type, "Deads"): run.data.metrics.get("DeadNeurons/test")
            }

            all_rows.append((row_key, metrics))

    # Build DataFrame
    records = {}
    for key, metrics in all_rows:
        if key not in records:
            records[key] = {}
        records[key].update(metrics)

    df = pd.DataFrame.from_dict(records, orient="index")
    df.index.names = ["Dimensions", "Expansion Ratio", "TopK"]

    # Calculate % differences and arrange columns
    result_cols = []
    for (dataset, metric) in sorted(set((k[0], k[2]) for k in df.columns)):
        main_col = (dataset, "main", metric)
        compare_col = (dataset, "compare", metric)
        percent_col = (dataset, "% change", metric)

        if main_col in df.columns and compare_col in df.columns:
            df[percent_col] = ((df[main_col] - df[compare_col]) / df[compare_col].abs()) * 100
            result_cols.extend([main_col, percent_col, compare_col])
        elif main_col in df.columns:
            result_cols.append(main_col)

    # Finalize view: sort columns, round values
    df = df[result_cols]
    df = df.round(4)
    df = df.sort_index(axis=1, level=[0, 2]).sort_values(
        by=["Dimensions", "Expansion Ratio", "TopK"],
        ascending=[True, True, True]
    )

    return df.reset_index()

sae_experiments = ['657713966175362303', '852893065079987597']

table = generate_sae_table_with_comparison(
    sae_experiments,
    main_note="sizes_Cosine",
    note_to_compare="sizes_L2",
)
table

table.mean()

experiments = ['523100174176986081', '333391697323445885']
table = generate_recommendations_with_comparision(
    experiments,
    aggregation_function="common_features",
    main_note="sizes_Cosine_without_acts",
    note_to_compare="sizes_L2_without_acts",
    group_type="sim",
)
table

table.groupby(["Dataset"]).mean()

table.groupby(["Dataset"]).std()

experiments = ['523100174176986081', '333391697323445885']
table = generate_recommendations_with_comparision(
    experiments,
    aggregation_function="common_features",
    main_note="sizes_Cosine_without_acts",
    note_to_compare="sizes_L2_without_acts",
    group_type="random",
)

table.groupby(["Dataset"]).mean()

experiments = ['523100174176986081', '333391697323445885']
table = generate_recommendations_with_comparision(
    experiments,
    aggregation_function="common_features",
    main_note="sizes_Cosine_without_acts",
    note_to_compare="sizes_L2_without_acts",
    group_type="div",
)
table

table.groupby(["Dataset"]).mean()

experiments = ['523100174176986081', '333391697323445885']
generate_recommendations_with_comparision(
    experiments,
    aggregation_function="average",
    main_note="sizes_Cosine_with_acts",
    note_to_compare="sizes_L2_with_acts",
    group_type="sim",
)

experiments = ['523100174176986081', '333391697323445885']
generate_recommendations_with_comparision(
    experiments,
    aggregation_function="average",
    main_note="sizes_Cosine_with_acts",
    note_to_compare="sizes_L2_with_acts",
    group_type="random",
)

experiments = ['523100174176986081', '333391697323445885']
generate_recommendations_with_comparision(
    experiments,
    aggregation_function="average",
    main_note="sizes_Cosine_with_acts",
    note_to_compare="sizes_L2_with_acts",
    group_type="div",
)